import pickle
import torch
import torchvision.transforms as T
import faiss
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import numpy as np
from PIL import Image

patch_h = 28
patch_w = 28
# 定义图像特征的维度
feat_dim = 384

transform = T.Compose([
    T.GaussianBlur(9, sigma=(0.1, 2.0)),
    T.Resize((patch_h * 14, patch_w * 14)),
    T.CenterCrop((patch_h * 14, patch_w * 14)),
    T.ToTensor(),
    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])

# 指定你想要搜索的目录
directory = './images/flowers'
# directory = 'E:\self-built-masked-face-recognition-dataset\AFDB_face_dataset'

images = []
# 过滤出图片文件
folders = [folder for folder in os.listdir(directory) if os.path.isdir(os.path.join(directory, folder))]

for folder in folders:
    folder_path = os.path.join(directory, folder)
    folder_files = os.listdir(folder_path)
    images.extend(
        [os.path.join(folder_path, file) for file in folder_files if file.endswith(".jpg") or file.endswith(".png")])

# 载入模型
dinov2 = torch.hub.load('./model', 'dinov2_vits14', source='local').cuda()

# 定义特征矩阵和图片URL映射表
feature_matrix = np.zeros((len(images), feat_dim), dtype=np.float32)
url_mapping = {}
# 创建Faiss索引
index = faiss.IndexFlatIP(feat_dim)

for i, image_path in enumerate(images):
    print("图片集载入中(" + str(i) + "/" + str(len(images)) + ")")
    image = Image.open(image_path)
    t_image = transform(image).unsqueeze(dim=0)
    faiss_feature = dinov2(t_image.cuda())
    normalized_feature = faiss_feature / faiss_feature.norm(dim=-1, keepdim=True)
    normalized_feature = normalized_feature.detach().cpu().numpy()
    # 将特征向量添加到特征矩阵中
    feature_matrix[i] = normalized_feature
    # 将特征向量与URL关联起来
    url_mapping[i] = image_path
    index.add(normalized_feature)

print("图片集载入向量库中...")
with open('./faiss/feature_matrix.pkl', 'wb') as f:
    pickle.dump(feature_matrix, f)

with open('./faiss/url_mapping.pkl', 'wb') as f:
    pickle.dump(url_mapping, f)

index = faiss.IndexFlatIP(feat_dim)

print("图片集载入向量库成功！")
